from pysurvival.models.semi_parametric import CoxPHModel
from pysurvival.utils.display import integrated_brier_score
from pysurvival.utils.metrics import concordance_index
from pysurvival.models.semi_parametric import NonLinearCoxPHModel
import pandas as pd

if __name__ == '__main__':
    data_path = 'data/dataset-b-day-i-cut.csv'
    origin_data = pd.read_csv(data_path)

    train_data = origin_data[:340]
    test_data = origin_data[340:]
    time_column = 'label'
    event_column = 'event'
    features = origin_data.columns.tolist()
    features.remove(time_column)
    features.remove(event_column)

    X_train, X_test = train_data[features], test_data[features]
    T_train, T_test = train_data[time_column], test_data[time_column]
    E_train, E_test = train_data[event_column], test_data[event_column]

    coxph_model = CoxPHModel()
    coxph_model.fit(X=X_train, T=T_train, E=E_train, init_method='he_uniforms',
                    l2_reg=0.1, lr=.5, tol=1e-4)

    c_index1 = concordance_index(model=coxph_model, X=X_test, T=T_test, E=E_test)
    print("CoxPH model c-index = {:.2f}".format(c_index1))

    ibs = integrated_brier_score(coxph_model, X_test, T_test, E_test, t_max=4,
                                 figure_size=(20, 6.5))
    print('IBS: {:.2f}'.format(ibs))

